"""
本代码是瞳孔识别的核心代码。
调用该代码后，会弹出两个窗口，左窗口可以动态调节瞳孔的阈值，右窗口展示瞳孔识别的效果，如果调整到合适的阈值，便可以在右窗口中用红色圆标识你的两个瞳孔。
若要使用，请在你的项目中保持本代码，并在你的代码中开头中直接from thresh_stack import stack_thresh
"""
import cv2
import dlib
import numpy as np
import cvzone
import tkinter as tk

def stack_thresh(cap):

    class PreviousValue:
        def __init__(self):
            self.xvalue = None
            self.yvalue = None

        def update(self, new_x_value, new_y_value):
            self.xvalue = new_x_value
            self.yvalue = new_y_value

        def xget(self):
            return self.xvalue

        def yget(self):
            return self.yvalue

    def shape_to_np(shape, dtype="int"):
        # initialize the list of (x, y)-coordinates
        coords = np.zeros((68, 2), dtype=dtype)
        # loop over the 68 facial landmarks and convert them
        # to a 2-tuple of (x, y)-coordinates
        for i in range(0, 68):
            coords[i] = (shape.part(i).x, shape.part(i).y)
        # return the list of (x, y)-coordinates
        return coords


    def eye_on_mask(mask, side):
        points = [shape[i] for i in side]
        points = np.array(points, dtype=np.int32)
        mask = cv2.fillConvexPoly(mask, points, 255)
        return mask


    def contouring(thresh, mid, img, right=False):
        cnts, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        cx, cy = 0, 0
        try:
            cnt = max(cnts, key=cv2.contourArea)
            M = cv2.moments(cnt)
            cx = int(M['m10'] / M['m00'])
            cy = int(M['m01'] / M['m00'])
            if right:
                cx += mid
            cv2.circle(img, (cx, cy), 4, (0, 0, 255), 2)
            return cx, cy
        except:
            return cx, cy
            pass

    detector = dlib.get_frontal_face_detector()
    predictor = dlib.shape_predictor('shape_predictor_68_face_landmarks.dat')

    left = [36, 37, 38, 39, 40, 41]
    right = [42, 43, 44, 45, 46, 47]

    ret, img = cap.read()
    thresh = img.copy()

    cv2.namedWindow('image')
    kernel = np.ones((9, 9), np.uint8)

    def nothing(x):
        pass

    previous_value = PreviousValue()
    cv2.createTrackbar('threshold', 'image', 0, 255, nothing)
    dis_thresh = 0

    while (True):
        # 读取人脸图像，进行灰度处理
        ret, img = cap.read()
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        rects = detector(gray, 1)
        for rect in rects:
            # 使用 NumPy 创建了一个新的黑色蒙版，其尺寸与网络摄像头框架相同。从关键点数组形状中存储左右眼睛点的XY坐标，并在蒙版上进行绘制
            shape = predictor(gray, rect)
            left_lx = shape.part(36).x
            left_ly = shape.part(36).y
            left_rx = shape.part(39).x
            left_ry = shape.part(39).y
            shape = shape_to_np(shape)
            mask = np.zeros(img.shape[:2], dtype=np.uint8)
            mask = eye_on_mask(mask, left)
            mask = eye_on_mask(mask, right)
            mask = cv2.dilate(mask, kernel, 5)

            # 在蒙版中绘制眼睛，并使用关键点 39 和 42 查找眼睛中点。通过对中点进行排序，找到中点两侧的最大等值线。可以利用它来找到眼球的中心。
            eyes = cv2.bitwise_and(img, img, mask=mask)
            mask = (eyes == [0, 0, 0]).all(axis=2)
            eyes[mask] = [255, 255, 255]
            mid = (shape[42][0] + shape[39][0]) // 2
            eyes_gray = cv2.cvtColor(eyes, cv2.COLOR_BGR2GRAY) # 将结果转换为灰度，使图像准备好进行阈值设置

            # 制作一个可调节的跟踪栏来控制阈值，从而能够更好的寻找到瞳孔
            threshold = cv2.getTrackbarPos('threshold', 'image')
            _, thresh = cv2.threshold(eyes_gray, threshold, 255, cv2.THRESH_BINARY)
            thresh = cv2.erode(thresh, None, iterations=2)  # 1
            thresh = cv2.dilate(thresh, None, iterations=4)  # 2
            thresh = cv2.medianBlur(thresh, 3)  # 3
            thresh = cv2.bitwise_not(thresh)

            # 获取左眼和右眼瞳孔中心的XY坐标
            left_cx, left_cy = contouring(thresh[:, 0:mid], mid, img)
            right_cx, right_cy = contouring(thresh[:, mid:], mid, img, True)

        cvzone.putTextRect(img, f'LEFTEYES_XY:{left_cx,left_cy}',  # 窗口显示内容是字符串类型，显示左瞳孔的XY坐标
                           (0, 70), 2, 1)

        cv2.imshow('eyes', img)
        cv2.imshow('image', thresh)

        if cv2.waitKey(1) & int(cv2.getWindowProperty('eyes', cv2.WND_PROP_VISIBLE)) == 0 & int(cv2.getWindowProperty('image', cv2.WND_PROP_VISIBLE)) == 0:
            break
    cap.release()